const fs = require('fs');
const zlib = require('zlib');
const { Readable } = require('stream');

// 将图像数据写入 BMP 文件
function writeImageToBmp(image, outputPath) {
  return new Promise((resolve, reject) => {
    const width = image.width;
    const height = image.height;
    const pixelData = image.data;

    // 计算每行的字节数，需要按 4 字节对齐
    const rowSize = Math.floor((width * 8 + 31) / 32) * 4;
    const dataSize = rowSize * height;

    // 创建 BMP 文件头
    const fileHeader = Buffer.alloc(14);
    fileHeader.write('BM', 0); // 文件类型
    fileHeader.writeUInt32LE(14 + 40 + dataSize, 2); // 文件大小
    fileHeader.writeUInt16LE(0, 6); // 保留字段
    fileHeader.writeUInt16LE(0, 8); // 保留字段
    fileHeader.writeUInt32LE(14 + 40, 10); // 数据偏移量

    // 创建信息头
    const infoHeader = Buffer.alloc(40);
    infoHeader.writeUInt32LE(40, 0); // 信息头大小
    infoHeader.writeInt32LE(width, 4); // 图像宽度
    infoHeader.writeInt32LE(height, 8); // 图像高度
    infoHeader.writeUInt16LE(1, 12); // 平面数
    infoHeader.writeUInt16LE(8, 14); // 位深度
    infoHeader.writeUInt32LE(0, 16); // 压缩方式
    infoHeader.writeUInt32LE(dataSize, 20); // 图像数据大小
    infoHeader.writeInt32LE(0, 24); // 水平分辨率
    infoHeader.writeInt32LE(0, 28); // 垂直分辨率
    infoHeader.writeUInt32LE(256, 32); // 调色板颜色数
    infoHeader.writeUInt32LE(0, 36); // 重要颜色数

    // 创建调色板
    const palette = Buffer.alloc(256 * 4);
    for (let i = 0; i < 256; i++) {
      palette.writeUInt8(i, i * 4);
      palette.writeUInt8(i, i * 4 + 1);
      palette.writeUInt8(i, i * 4 + 2);
      palette.writeUInt8(0, i * 4 + 3);
    }

    // 创建图像数据缓冲区
    const imageData = Buffer.alloc(dataSize);
    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        const srcIndex = (height - y - 1) * width + x;
        const dstIndex = y * rowSize + x;
        imageData.writeUInt8(pixelData[srcIndex], dstIndex);
      }
    }

    // 合并所有部分
    const bmpData = Buffer.concat([fileHeader, infoHeader, palette, imageData]);

    // 写入文件
    fs.writeFile(outputPath, bmpData, (err) => {
      if (err) {
        reject(err);
      } else {
        resolve();
      }
    });
  });
}

async function testWriteImages(images, count) {
  for (let i = 0; i < count; i++) {
    const outputPath = `image_${i}.bmp`;
    await writeImageToBmp(images[i], outputPath);
    console.log(`已保存图像到 ${outputPath}`);
  }
}

// 直接解压并处理数据
function getUbyteFileData(inputPath) {
  return new Promise((resolve, reject) => {
    const readStream = fs.createReadStream(inputPath);
    const gunzip = zlib.createGunzip();
    const chunks = [];
    readStream.pipe(gunzip).on('data', (chunk) => { // 解压.gz文件
      chunks.push(chunk);
    }).on('end', () => {
      const buffer = Buffer.concat(chunks);
      const readableStream = new Readable();
      readableStream.push(buffer);
      readableStream.push(null);
      parseUbyteStream(readableStream)
        .then((data) => {
          resolve(data);
        })
        .catch((err) => {
          reject(err);
        });
    }).on('error', (err) => {
      reject(err);
    });
  });
}

// 解析流中的 .ubyte 数据
function parseUbyteStream(stream) {
  return new Promise((resolve, reject) => {
    const chunks = [];
    stream.on('data', (chunk) => {
      chunks.push(chunk);
    }).on('end', () => {
      const data = Buffer.concat(chunks);
      const magicNumber = data.readUInt32BE(0); // 文件类型
      const numItems = data.readUInt32BE(4); // 数据项数量

      let images = [];
      if (magicNumber === 2051) { // 图像文件
        const rows = data.readUInt32BE(8);
        const cols = data.readUInt32BE(12);
        const imageSize = rows * cols;

        for (let i = 0; i < numItems; i++) {
          const startIndex = 16 + i * imageSize;
          const endIndex = startIndex + imageSize;
          const imageData = data.slice(startIndex, endIndex);
          const image = [];
          for (let j = 0; j < imageSize; j++) {
            image.push(imageData.readUInt8(j));
          }
          images.push({ width: cols, height: rows, data: image });
        }
      } else if (magicNumber === 2049) { // 标签文件
        for (let i = 0; i < numItems; i++) {
          const label = data.readUInt8(8 + i);
          images.push(label);
        }
      }
      resolve(images);
    }).on('error', (err) => {
      reject(err);
    });
  });
}

module.exports = {
  getUbyteFileData,
}

// 主函数
async function test() {
  try {
    // 解压并解析文件
    const labels = await getUbyteFileData('./MINIST/train-labels-idx1-ubyte.gz');
    console.log('解析结果:', labels.length);
    // const images = await getUbyteFileData('./MINIST/train-images-idx3-ubyte.gz');
    // testWriteImages(images, 2);
    // console.log('=================images[0]', images[0]);
  } catch (error) {
    console.error('处理过程中出现错误:', error);
  }
}

if (require.main === module) {
  test();
}
